import tensorflow as tf

# from tensorflow import keras

x = tf.random.normal([10, 784])
net = tf.keras.layers.Dense(512)
out = net(x)

print(x.shape)
# print(out)
print(out.shape)
print(net.kernel.shape)
print(net.bias.shape)

model = tf.keras.Sequential([
    tf.keras.layers.Dense(512, activation='relu'),
    tf.keras.layers.Dense(256, activation='relu'),
    tf.keras.layers.Dense(10, activation='relu')
])

model.build(input_shape=[None, 784])
model.summary()

for p in model.trainable_variables:
    # print(p)
    print(p.name, p.shape)
